import random
import torch

from torchvision.transforms import functional as F


def _flip_coco_person_keypoints(kps, width):
    """翻转COCO格式的人体关键点坐标
    Args:
        kps: 原始关键点坐标 [num_persons, 17, 3] (x, y, visibility)
        width: 图像宽度
    Returns:
        flipped_data: 水平翻转后的关键点坐标
    """
    # 关键点索引映射关系（左右对称点交换）
    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
    flipped_data = kps[:, flip_inds]
    # 水平翻转x坐标
    flipped_data[..., 0] = width - flipped_data[..., 0]
    # 处理不可见点（visibility=0时坐标归零）
    inds = flipped_data[..., 2] == 0
    flipped_data[inds] = 0
    return flipped_data


class Compose(object):
    """组合多个数据增强变换（按顺序应用）"""
    def __init__(self, transforms):
        """
        Args:
            transforms (list): 包含多个数据增强操作的列表
        """
        self.transforms = transforms

    def __call__(self, image, target):
        """应用所有变换
        Args:
            image: 输入图像
            target: 包含标注信息的字典
        Returns:
            (image, target): 处理后的图像和标注
        """
        for t in self.transforms:
            image, target = t(image, target)
        return image, target


class RandomHorizontalFlip(object):
    """以给定概率随机水平翻转图像及标注"""
    def __init__(self, prob):
        """
        Args:
            prob (float): 翻转概率，范围[0,1]
        """
        self.prob = prob

    def __call__(self, image, target):
        """执行随机翻转操作
        Args:
            image: 输入图像张量
            target: 包含标注的字典
        Returns:
            (image, target): 可能翻转后的图像和标注
        """
        if random.random() < self.prob:
            height, width = image.shape[-2:]
            image = image.flip(-1)  # 水平翻转图像
            # 处理边界框坐标
            bbox = target["boxes"]
            bbox[:, [0, 2]] = width - bbox[:, [2, 0]]
            target["boxes"] = bbox
            # 处理掩码（如果有）
            if "masks" in target:
                target["masks"] = target["masks"].flip(-1)
            # 处理关键点（如果有）
            if "keypoints" in target:
                keypoints = target["keypoints"]
                keypoints = _flip_coco_person_keypoints(keypoints, width)
                target["keypoints"] = keypoints
        return image, target


class ToTensor(object):
    """将PIL图像或numpy数组转换为张量"""
    def __call__(self, image, target):
        """
        Args:
            image: PIL图像或numpy数组
            target: 标注字典
        Returns:
            (image_tensor, target): 转换后的张量和原标注
        """
        image = F.to_tensor(image)  # 转换为[0,1]范围的张量
        return image, target
